package com.spark.sniffer.plugin.plugins.interceptor;

import java.lang.reflect.Method;

import com.spark.sniffer.logging.api.Logger;
import com.spark.sniffer.logging.api.LoggerFactory;
import com.spark.sniffer.plugin.interceptor.enhance.InstanceMethodsAroundInterceptor;
import com.spark.sniffer.plugin.interceptor.enhance.MethodInterceptResult;
import com.spark.sniffer.tracing.TracingContext;

/**
 * 拦截Connection.prepareStatement方法，将traceid注入SQL语句中
 * 
 * @author cobe
 * @date 2021/07/21
 */
public class JdbcConnectionInterceptor implements InstanceMethodsAroundInterceptor {

    private static final Logger log = LoggerFactory.getLogger(JdbcConnectionInterceptor.class);

    @Override
    public void beforeMethod(Object instance, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
        MethodInterceptResult result) throws Throwable {
        if (allArguments == null || allArguments.length == 0 || !(allArguments[0] instanceof String)
            || TracingContext.getContext().getTraceId() == null) {
            return;
        }

        String sql = (String)allArguments[0];
        String injectTraceId =
            new StringBuilder(" /* ").append(TracingContext.getContext().getTraceId()).append(" */ ").toString();
        if (sql.contains(injectTraceId)) {
            return;
        }

        injectTraceId(allArguments, injectTraceId);
    }

    @Override
    public Object afterMethod(Object instance, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
        Object ret) throws Throwable {
        return ret;
    }

    @Override
    public void handleMethodException(Object instance, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
        Throwable t) {}

    private void injectTraceId(Object[] allArguments, String injectTraceId) {
        String sql = ((String)allArguments[0]).trim();

        int firstSpaceIndex = sql.indexOf(" ");
        sql = new StringBuilder(sql.substring(0, firstSpaceIndex).trim()).append(injectTraceId)
            .append(sql.substring(firstSpaceIndex).trim()).toString();
        allArguments[0] = sql;

        log.debug((String)allArguments[0]);
    }
}
